#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
读取和分析组织化的脑电数据段演示

支持的数据格式：
1. 端口模式: {通道名: [seg1, seg2, ...]} for each channel
2. 语料范式: {动作X_准备: [repeat1, repeat2, ...], 动作X_开始: [repeat1, repeat2, ...], 动作X_复位: [repeat1, repeat2, ...], 通道4: [seg1, seg2, ...], ...}
3. 运动范式: {动作X_准备: [repeat1, repeat2, ...], 动作X_开始: [repeat1, repeat2, ...], 动作X_复位: [repeat1, repeat2, ...], 通道4: [seg1, seg2, ...], ...}

支持的动作阶段：动作1_准备, 动作1_开始, 动作1_复位, 动作2_准备, 动作2_开始, 动作2_复位, ...
特殊通道：通道4
"""

import numpy as np
import os

def load_organized_segments(file_path):
    """
    加载组织好的脑电数据段
    
    Args:
        file_path (str): 数据段文件路径
    
    Returns:
        dict: 组织好的脑电数据段
    """
    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        return None
    
    # 加载数据
    data = np.load(file_path, allow_pickle=True).item()
    
    print(f"=== 加载脑电数据段 ===")
    print(f"文件: {os.path.basename(file_path)}")
    print(f"原始文件: {data.get('original_file', 'N/A')}")
    print(f"采样频率: {data.get('fs', 'N/A')} Hz")
    print(f"数据类型: {data.get('data_type', 'port')}")
    
    # 显示音频信息
    if 'audio_segments' in data and data['audio_segments']:
        print(f"原始音频文件: {data.get('original_audio_file', 'N/A')}")
        print(f"音频采样率: {data.get('audio_fs', 'N/A')} Hz")
        print(f"音频通道数: {data.get('audio_channels', 'N/A')}")
        print(f"音频分割段数: {len(data['audio_segments'])} (仅基于通道1标记)")
    else:
        print("无音频数据")
    
    if data.get('data_type') == 'corpus':
        print(f"范式CSV文件: {data.get('corpus_csv_path', 'N/A')}")
        corpus_data = data.get('corpus_data', [])
        if corpus_data:
            # 检测是运动范式还是语料范式
            if 'channel1_text' in corpus_data[0] and 'channel2_text' in corpus_data[0]:
                # 通过检查CSV路径和内容判断范式类型
                csv_path = data.get('corpus_csv_path', '')
                if 'condition.csv' in csv_path.lower():
                    print("范式类型: 运动范式 (condition.csv)")
                    print("范式信息:")
                    for item in corpus_data:
                        print(f"  {item['corpus_order']}: 通道1={item['channel1_text']}, 通道2={item['channel2_text']}")
                else:
                    print("范式类型: 语料范式 (welinking_corpus.csv)")
                    print("语料信息:")
                    for item in corpus_data:
                        print(f"  {item['corpus_order']}: {item['channel1_text']}")
            else:
                # 兼容旧格式
                print("语料信息:")
                for item in corpus_data:
                    if 'wl_corpus' in item:
                        print(f"  {item['corpus_order']}: {item['wl_corpus']}")
                    else:
                        print(f"  {item['corpus_order']}: {item}")
    else:
        print(f"端口配置: {data.get('port_durations', 'N/A')}")
    
    grouped_data = data.get('grouped_data', {})
    data_type = data.get('data_type', 'port')
    
    if data_type == 'corpus':
        print(f"\n=== 范式数据结构 ===")
        for key, segments in grouped_data.items():
            if key.startswith('动作'):
                # 动作阶段数据：显示为重复次数
                print(f"动作阶段'{key}':")
                print(f"  重复次数: {len(segments)}")
                
                for i, seg in enumerate(segments):
                    print(f"  重复{i+1}: 形状{seg.shape} (通道数x样本数)")
                    if i >= 2:  # 只显示前3段
                        print(f"  ... 还有{len(segments)-3}次重复")
                        break
            else:
                # 特殊通道：显示为数据段
                print(f"特殊通道'{key}':")
                print(f"  数据段数量: {len(segments)}")
                
                for i, seg in enumerate(segments):
                    print(f"  段{i+1}: 形状{seg.shape} (通道数x样本数)")
                    if i >= 2:  # 只显示前3段
                        print(f"  ... 还有{len(segments)-3}段")
                        break
    else:
        print(f"\n=== 端口数据结构 ===")
        for port, segments in grouped_data.items():
            print(f"{port}:")
            print(f"  数据段数量: {len(segments)}")
            
            for i, seg in enumerate(segments):
                print(f"  段{i+1}: 形状{seg.shape} (通道数x样本数)")
                if i >= 2:  # 只显示前3段
                    print(f"  ... 还有{len(segments)-3}段")
                    break
    
    return data


def analyze_segments_data(data):
    """
    分析脑电数据段
    
    Args:
        data (dict): 加载的数据
    """
    if not data or 'grouped_data' not in data:
        print("数据格式错误")
        return
    
    grouped_data = data['grouped_data']
    
    print(f"\n=== 数据分析 ===")
    
    # 统计每个端口/范式的数据
    total_segments = 0
    data_type = data.get('data_type', 'port')
    
    for key, segments in grouped_data.items():
        total_segments += len(segments)
        
        # 计算每段的平均长度
        lengths = [seg.shape[1] for seg in segments]
        avg_length = np.mean(lengths)
        
        if data_type == 'corpus':
            if key.startswith('动作'):
                print(f"动作阶段'{key}':")
                print(f"  重复次数: {len(segments)}")
                print(f"  平均长度: {avg_length:.1f} 样本")
                print(f"  总样本数: {sum(lengths)}")
            else:
                print(f"特殊通道'{key}':")
                print(f"  段数: {len(segments)}")
                print(f"  平均长度: {avg_length:.1f} 样本")
                print(f"  总样本数: {sum(lengths)}")
        else:
            print(f"{key}:")
            print(f"  段数: {len(segments)}")
            print(f"  平均长度: {avg_length:.1f} 样本")
            print(f"  总样本数: {sum(lengths)}")
        
        if len(segments) > 0:
            print(f"  通道数: {segments[0].shape[0]}")
    
    print(f"\n总数据段: {total_segments}")
    
    # 显示音频分割信息
    if 'audio_segments' in data and data['audio_segments']:
        audio_segments = data['audio_segments']
        audio_fs = data.get('audio_fs', 44100)
        print(f"\n=== 音频分割分析 ===")
        print(f"说明: 音频分割仅基于通道1的时间标记，与CSV顺序一致")
        
        for key, segments_list in audio_segments.items():
            total_duration = sum(len(seg) / audio_fs for seg in segments_list)
            print(f"音频分组'{key}': {len(segments_list)}个片段, 总时长{total_duration:.3f}s")
            
            for idx, segment in enumerate(segments_list):
                duration = len(segment) / audio_fs
                print(f"  片段{idx+1}: {duration:.3f}s, 形状{segment.shape}")
                if idx >= 2:  # 只显示前3个片段
                    print(f"  ... 还有{len(segments_list)-3}个片段")
                    break


def get_grouped_data(data, key):
    """
    获取指定端口、动作阶段或特殊通道的所有数据段
    
    Args:
        data (dict): 加载的数据
        key (str): 端口名称、动作阶段名称或特殊通道名称
                  动作阶段格式：'动作1_准备', '动作2_开始', '动作1_复位' 等
                  特殊通道：'通道4' 等
    
    Returns:
        list: 该端口/范式/特殊通道的所有数据段
    """
    grouped_data = data.get('grouped_data', {})
    data_type = data.get('data_type', 'port')
    
    if key not in grouped_data:
        if data_type == 'corpus':
            if key.startswith('动作'):
                print(f"动作阶段'{key}'没有数据")
            else:
                print(f"特殊通道'{key}'没有数据")
        else:
            print(f"{key}没有数据")
        return []
    
    segments = grouped_data[key]
    if data_type == 'corpus':
        if key.startswith('动作'):
            print(f"动作阶段'{key}'有{len(segments)}次重复")
        else:
            print(f"特殊通道'{key}'有{len(segments)}个数据段")
    else:
        print(f"{key}有{len(segments)}个数据段")
    
    return segments


def concatenate_grouped_data(data, key):
    """
    将指定端口、动作阶段或特殊通道的所有段连接成一个大的数组
    
    Args:
        data (dict): 加载的数据
        key (str): 端口名称、动作阶段名称或特殊通道名称
    
    Returns:
        np.ndarray: 连接后的数据
    """
    segments = get_grouped_data(data, key)
    
    if not segments:
        return None
    
    # 连接所有段
    concatenated = np.concatenate(segments, axis=1)
    data_type = data.get('data_type', 'port')
    
    if data_type == 'corpus':
        if key.startswith('动作'):
            print(f"动作阶段'{key}'连接后的数据形状: {concatenated.shape}")
        else:
            print(f"特殊通道'{key}'连接后的数据形状: {concatenated.shape}")
    else:
        print(f"{key}连接后的数据形状: {concatenated.shape}")
    
    return concatenated


def main():
    """主函数"""
    # 要求用户输入文件路径
    while True:
        segments_file = input("请输入数据段文件路径: ").strip()
        
        if not segments_file:
            print("路径不能为空，请重新输入")
            continue
        
        if not segments_file.endswith('.npy'):
            print("请输入.npy文件")
            continue
        
        if os.path.exists(segments_file):
            break
        else:
            print(f"文件不存在: {segments_file}")
            print("请检查路径是否正确")
    
    print(f"使用文件: {segments_file}")
    
    # 加载数据
    data = load_organized_segments(segments_file)
    
    if data:
        # 分析数据
        analyze_segments_data(data)
        
        # 演示如何使用数据
        grouped_data = data.get('grouped_data', {})
        
        if grouped_data:
            # 获取第一个端口的数据
            first_port = list(grouped_data.keys())[0]
            print(f"\n=== 演示: 使用{first_port}的数据 ===")
            
            # 获取第一段数据
            segments = get_grouped_data(data, first_port)
            if segments:
                first_segment = segments[0]
                print(f"第一段数据形状: {first_segment.shape}")
                print(f"数据类型: {first_segment.dtype}")
                print(f"数据范围: {first_segment.min():.2f} 到 {first_segment.max():.2f}")
                
                # 连接所有段
                concatenated = concatenate_grouped_data(data, first_port)
        
        print(f"\n=== 使用说明 ===")
        data_type = data.get('data_type', 'port')
        if data_type == 'corpus':
            # 检测是运动范式还是语料范式
            csv_path = data.get('corpus_csv_path', '')
            if 'condition.csv' in csv_path.lower():
                print(f"运动范式数据结构:")
                print(f"  脑电数据: data['grouped_data']['动作X_准备'][重复索引]")
                print(f"  音频数据: data['audio_segments']['动作X_准备'][重复索引]")
                print(f"")
                print(f"示例:")
                print(f"  获取'动作1_准备'的第2次脑电重复: data['grouped_data']['动作1_准备'][1]")
                print(f"  获取'动作1_准备'的第2次音频重复: data['audio_segments']['动作1_准备'][1]")
            else:
                print(f"语料范式数据结构:")
                print(f"  脑电数据: data['grouped_data']['语料001_XXX_01'][段索引]")
                print(f"  音频数据: data['audio_segments']['语料001_XXX_01'][段索引]")
                print(f"")
                print(f"示例:")
                print(f"  获取第1个语料的脑电: data['grouped_data']['语料001_你_01'][0]")
                print(f"  获取第1个语料的音频: data['audio_segments']['语料001_你_01'][0]")
            print(f"每段数据是numpy数组")
            print(f"  脑电数据形状: (通道数, 样本数)")
            print(f"  音频数据形状: (样本数,) 或 (样本数, 音频通道数)")
        else:
            print(f"端口模式数据结构:")
            print(f"  脑电数据: data['grouped_data'][通道名][段索引]")
            print(f"  音频数据: data['audio_segments'][通道名][段索引]")
            print(f"")
            print(f"示例:")
            print(f"  获取通道1的第2段脑电: data['grouped_data']['通道1'][1]")
            print(f"  获取通道1的第2段音频: data['audio_segments']['通道1'][1]")
            print(f"每段数据是numpy数组")
            print(f"  脑电数据形状: (通道数, 样本数)")
            print(f"  音频数据形状: (样本数,) 或 (样本数, 音频通道数)")
        
        # 音频保存提示
        if 'audio_segments' in data and data['audio_segments']:
            print(f"\n=== 音频保存 ===")
            print(f"使用 save_audio_segments.py 脚本可将音频段保存为wav文件:")
            print(f"  python save_audio_segments.py {os.path.basename(segments_file)}")
            print(f"这将创建 audio_segments/ 目录并保存所有音频分割段")
            print(f"\n音频分割说明:")
            print(f"  - 音频分割仅基于通道1的时间标记")
            print(f"  - 分割顺序严格按照CSV中的order字段顺序")
            print(f"  - 第1段音频对应第1个CSV条目，第2段对应第2个，以此类推")


if __name__ == "__main__":
    main()